'''
Created on Nov 29, 2010

:author: rreyes
'''

from yacf.Backends.Common.Visitors.GenericVisitors import UndefIDFilter,\
    FuncCallFilter, FuncDeclOfNameFilter

from yacf.Tools.Tree import NodeNotFound, getFileNode
from yacf.Frontend.Shortcuts import getCurrentLanguageAst
import config
from yacf.Frontend.C99 import c99_ast
c_ast = getCurrentLanguageAst()

from yacf.Backends.Common.Mutators.AbstractMutator import AbstractMutator
from yacf.Backends.Common.Visitors.GenericVisitors import GenericFilterVisitor
from yacf.Tools.Tree import NodeNotFound, NodeNotValid
from yacf.Frontend.SymbolTable import getSymbolTableFromNode
from yacf.Backends.Common.TemplateEngine.TemplateParser import TemplateParser, get_template_array
from yacf.Frontend.InternalRepr import AstToIR
from yacf.Tools.Tree import ReplaceTool, getContainerAttribute, getFileNode
import copy
from yacf.Backends.Common.Visitors.GenericVisitors import FuncDefOfNameFilter,\
    FuncCallFilter
from yacf.Tools.Tree import getFileNode, NodeNotFound




class InstructionFlowBreakFilter(GenericFilterVisitor):
    """ Return the first matching flow break statement
    """   
    def __init__(self, prev_brother = None):
        FLOW_BREAK_STMTS = [c_ast.Return, c_ast.Break, c_ast.Goto]
        self._isNested = False
        
        def condition(node):
            if type(node) in [c_ast.Return, c_ast.Goto]:
                return True
            elif type(node) in [c_ast.Break] and not self._isNested:
                return True
            return False
                
        super(InstructionFlowBreakFilter, self).__init__(condition_func = condition, prev_brother = prev_brother)
        
    
    def visit_For(self, node, offset, ignore = []):
        self._isNested = True
        self.generic_visit(node.init, offset,)
        self.generic_visit(node.cond, offset,)
        self.generic_visit(node.next, offset,)
        self.generic_visit(node.stmt, offset,)
        self.isNested = False
    

class DeclarationOrUseOfPointersFilter(GenericFilterVisitor):
    """ Return the first declaration or use of a pointer
    """   
    def __init__(self, prev_brother = None):
        
        def condition(node):
            if type(node) == c_ast.PtrDecl:
                return True
            elif type(node) == c_ast.UnaryOp and node.op == '*':
                return True
            return False
                
        super(DeclarationOrUseOfPointersFilter, self).__init__(condition_func = condition, prev_brother = prev_brother)
      
def compound_has_got_pointers(block_node):
    """ Return true if exists the declaration or use of a pointer in the block node
    
        :param block_node: The group of nodes where the function will apply the filter
    """
    try:
        DeclarationOrUseOfPointersFilter().apply(block_node)
        return True
    except NodeNotFound:
        return False

def block_is_SESE(block_node):
    try:
        InstructionFlowBreakFilter().apply(block_node)
    except NodeNotFound:
        return True
    raise NodeNotValid("Not a SESE Block")


class ParameterStyle:
        C_POINTER = '*'
        CXX_REFERENCE = '&'

class OutlineMutator(AbstractMutator):
    """ Outline a code block into a separate function.
    
    Given the subtree reference to a SESE block (called *sese_block* from now on), the Outliner involves:

    1. Creating a function *outlined_function* whose body is a copy of *sese_block*
    2. Replace the appropriate occurrences of C in the original AST by a function call to *outlined_function*.

    Outlining might be seen as a trivial task, but some consideration must be taken into account.
    Special attention must be given to the type definitions of the variables used inside the block *sese_block*. 
    
    The outlining algorithm is language independent (as far as the language has support for subroutines), but 
    there are several syntactic details involved in the generated code that it is not possible to decouple the
    algorithm from the implementation and therefore, a Strategy pattern cannot be used.
    
    To leverage the language-dependency, two templates have been defined for the translation: _outline_function_template and
    _outline_call_template. The first is in charge of constructing the code related to the destination function, and the latter
    is devoted to the function call creation.
    
    This class is heavily used within the |yacf| framework and therefore some hooks have been added to increase the
    flexibility. This hooks are: 
    * Defining the outlined function name, by setting the name parameter at the constructor 
    * Applying the outline or not: Sometimes, the developer don't want to replace the original code, 
        but she wants to obtain the outlined function. Setting replace parameter to False will change
        the default behavior of the mutator by doing all except of the replacing. 
    
    """
    

    
    
    def __init__(self, name = "foo", replace = True, 
                 parameter_style = ParameterStyle.C_POINTER, 
                 outline_function_template = None, 
                 outline_call_template = None, change_code_block = True):
        """
            :param name: Name of the outlined function.
            :param boolean replace: Specify if *sese_block* will be replaced by the function call or not.
            :param ParameterStyle parameter_style: Parameter passing style to be used in the outlined function
        """
        self._name = name
        self.outlined_subtree = None
        self._replace = replace
        self._extra_decls = []
        self.prototype_subtree = None
        self._parameter_style = parameter_style;
        self._parameters = []
        self._change_code_block = change_code_block
        self._undeclFunctions = []
        if not outline_function_template:
            if self._parameter_style == ParameterStyle.C_POINTER:
                self._outline_function_template = """
                 <%!
                       from yacf.Backends.Common.TemplateEngine.Functions import is_complex_type, is_array, is_struct, is_pointer
                 %>

		            %for var in typedef:
		               ${var};
		            %endfor
		
		            %for var in set(variables):
                        %if var.structDecl != "":
		                    ${var.structDecl}
                        %endif
		            %endfor
		
                    <%
                       my_parameter_array = []
                    %>
                           
                    %for var in parameters:
                      %if is_array(var.ptr):
                         <%
                           #if is_array(var.ptr.type):
                               my_parameter_array.append(str(var.declaration) + "*" + str(var))
                           #else:
                           #    my_parameter_array.append(str(var.declaration) + str(var))
                         %>
                      %elif is_pointer(var.ptr):
                         <%
                           my_parameter_array.append(str(var.declaration) + str(var))
                         %>
                      %else:
                         <%
                           my_parameter_array.append(str(var.declaration) + "*_" + str(var))
                         %>
                      %endif
                    %endfor

		              
		            void ${funcName} (
                        ${', '.join(str(var) for var in my_parameter_array)}    
		                );
		                
		            void ${funcName} (
                        ${', '.join(str(var) for var in my_parameter_array)}    
		            ) {
		                %for var in parameters:
                            %if not is_complex_type(var.ptr):
		                        ${str(var.declaration) + str(var) + " = " +  "*_" + str(var)};
                            %endif
		                %endfor   
		                ${code};
		                %for var in parameters:
                            %if not is_complex_type(var.ptr):
    		                    ${"*_" + str(var) + " = " + str(var)};
                            %endif
		                %endfor   
		            }
                """
            else:
                self._outline_function_template = """
		            %for var in variables:
		                ${var.typedef or ""}
		            %endfor
		
		            %for var in variables:
		                ${var.structDecl or ""}
		            %endfor
		
		              
		            void ${funcName} (
		                ${', '.join(str(var.declaration) + str(var) for var in parameters)}    
		                );
		                
		            void ${funcName} (
		            ${', '.join(str(var.declaration) + str(var) for var in parameters)}    
		            ) {
		                ${code};
		            }
                """

        else:
            self._outline_function_template = outline_function_template
        if not outline_call_template:
            if self._parameter_style == ParameterStyle.C_POINTER:
                self._outline_call_template = """
                 <%!
                       from yacf.Backends.Common.TemplateEngine.Functions import is_complex_type
                 %>
		            %for var in typedef:
		               ${var};
		            %endfor
            
                    %for var in parameters:
                        ${var.structDecl or ""}
                    %endfor

            
                    int fake () {
                    ${funcName} (
                        <%
                            my_parameter_array = []
                        %>
                           
                        %for var in parameters:
                          %if not is_complex_type(var.ptr):
                            <%
                                my_parameter_array.append('&' + str(var))
                            %>
                          %else:
                            <%
                                my_parameter_array.append(str(var))
                            %>
                          %endif
                        %endfor

                        ${', '.join(str(var) for var in my_parameter_array)}    
                    );
                    }
                """
            else:
                self._outline_call_template = """
             %for var in parameters:
                ${var.typedef or ""}
            %endfor
            
            %for var in parameters:
                ${var.structDecl or ""}
            %endfor

            
            int fake () {
            ${funcName} (
            ${', '.join('&' + str(var) for var in parameters)}    
            );
            }
                """
        else:
            self._outline_call_template = outline_call_template
        self._typedefList = []
        self._funcCallList = []
        super(OutlineMutator, self).__init__()
    
    @property
    def outline_call_template(self):
        return self._outline_call_template
    
    @property
    def outline_function_template(self):
        return self._outline_function_template
    
    def filter(self, ast):
        return ast
    
    def getOutlinedFunction(self):
        return self.outlined_subtree  
    
    def setOutlinedFunction(self, subtree):
        self.outlined_subtree = subtree
    
    def getOutlinedPrototype(self):
        return self.prototype_subtree

    def getOutlinedFunctionCall(self):
        return self._funccall

    def getTypedefList(self):
        return self._typedefList
    
    def getFuncCallList(self):
        """ Return the list of function calls within a block.
            This method is mantained only for compatibility with
            old-fashion backends (like MPI and/or pure Cuda). It
            returns a list of function declarations of the functions
            called inside the outlined function.

            .. deprecated::
                Use funcCalls property instead. 
        """
        called_list = []
        undeclared_list = []
        for elem in self._funcCallList:
            try:
                funcDecl = FuncDeclOfNameFilter(name = elem.name).apply(getFileNode(self.main_ast))
            except NodeNotFound:
                # Declaration could not be found
                undeclared_list.append(elem)
                continue
            called_list.append(funcDecl.parent) 
        return called_list #, undeclared_list]
 
    @property
    def funcCalls(self):
        """ List of function calls inside the kernel """
        return self._funcCallList 


    def getParameterList(self):
        return self._parameters  
    


    def _declToParam (self, node):
        """ Create a param node from a variable declaration.

            :param node: Declaration to translate
            :return:  Parameter form of the declaration

        """
        param = node.__copy__()
        # Quick return if possible
        if isinstance(param.type, c_ast.PtrDecl):
            return param
        
        if self._parameter_style == ParameterStyle.C_POINTER:
            from yacf.Backends.Common.Mutators.AstSupport import PointerMutator
            PointerMutator().apply(param)
        elif self._parameter_style == ParameterStyle.CXX_REFERENCE:
            from yacf.Backends.Common.Mutators.AstSupport import ReferenceMutator
            ReferenceMutator().apply(param)

        return param
    
    def buildOutlineFunction(self,  org_decls, target):
        """ Build a tree with the outlined function """
        from yacf.Backends.Common.Visitors.GenericVisitors import IDFilter
        from yacf.Tools.Tree import getContainerAttribute
        from yacf.Tools.Tree import ReplaceTool

        decls = [ elem.__copy__() for elem in org_decls]
        # Transform the declaration list into a parameter declaration
        params = [ self._declToParam(elem) for elem in decls]
        # Change variable access to value of variable access
        st = getSymbolTableFromNode(elem)

        if self._change_code_block:
                for elem in decls:
                    nodes_to_change = []
                    # TODO: Put in one loop
                    for node in IDFilter(id = c_ast.ID(name = elem.name,)).dfs_iter(target):
                        #===============================================================
                        # import pdb
                        # pdb.set_trace()
                        #===============================================================
                        symbol = st.lookUp(node)
                        if not isinstance(symbol.type, c_ast.ArrayDecl, ) and not isinstance(symbol.type, c_ast.PtrDecl):
                            nodes_to_change.append(node)
                    
                    for id in nodes_to_change:
                        value_node = None
                        if self._parameter_style == ParameterStyle.C_POINTER:
                            value_node = c_ast.ExprList(exprs = [c_ast.UnaryOp(op = '*', expr = id,)])
                            value_node.exprs[0].parent = value_node
                        elif self._parameter_style == ParameterStyle.CXX_REFERENCE:
                            value_node = id
                        else:
                            raise Exception

                        ReplaceTool(new_node = value_node, old_node = id).apply(id.parent, getContainerAttribute(id))
                        from yacf.Frontend.SymbolTable import updateSymbolTableFromNode
                        updateSymbolTableFromNode(self.main_ast)

        
        #  print "****************************"

        self._parameters =  get_template_array(params, self.main_ast, st = st)
        subs_dir = {'variables' : get_template_array(decls, self.main_ast) , 
                        'parameters' : self._parameters, 
                        'code' : str(target), 
                        'funcName' : self._name,
                        'ArrayDecl' : str(c_ast.ArrayDecl),
                        'PtrDecl' : str(c_ast.PtrDecl),
                        'typedef': self.get_typeDefinitions(decls, self.main_ast),}
        template_code = TemplateParser(self.outline_function_template).render(**subs_dir)
        #  print "****************************"
        subtree = self.parse_snippet(template_code, subs_dir, 'fbody', show = False, languages = config.CURRENT_LANGUAGES)
        
        return [subtree.ext[-1], subtree.ext[-2]]

    def get_typeDefinitions(self, decls, main_ast):
        """ Build a list with all typedefs used in declarations or parameters.
        """
        declarations = get_template_array(decls, main_ast)
        typedef = list()
        for decl in declarations:
            if decl.typedef and not decl.typedef in typedef:
                typedef.append(decl.typedef)
        return typedef
        

    def add_additional_decls(self, decl_list):
        """ The declarations in the decl_list are added as extra parameters to the 
                outlined function. 
        """
        self._extra_decls += decl_list


    def buildFunctionTypedefList(self, params, target):
        """ Build a list with the typedefs required for the outlined function. """
        # Outline template
        subs_dir = {'parameters' : get_template_array(params, self.main_ast), 
                    'code' : str(target), 
                    'funcName' : self._name, 
                    'typedef': self.get_typeDefinitions(params, self.main_ast),}
        template_code = TemplateParser(self.outline_call_template).render(**subs_dir)
        return self.parse_snippet(template_code, subs_dir, 'tdef', show = False).ext[0:-1]
    
    def buildOutlineCall(self, params, target):
        """ Build a function call to the outlined function.
        
            :param params: Params of the function
            :return: Subtree with the function call """
        # Outline template
        subs_dir = {'parameters' : get_template_array(params, self.main_ast), 
                    'funcName' : self._name, 
                    'ArrayDecl': str(c_ast.ArrayDecl),
                    'typedef': self.get_typeDefinitions(params, self.main_ast),}
        template_code = TemplateParser(self.outline_call_template).render(**subs_dir)
        return self.parse_snippet(template_code, subs_dir, 'fcall', show = False).ext[-1].body.block_items[-1]
        
    
    def mutatorFunction(self, sese_block):
        """
             Apply the OutlineMutation to the given sese_block.
             The following class attributed are modified after the mutation:
             * _funccall
             * _typedefList
             * _funcCallList
             * outlined_subtree
             * prototype_subtree


             :param sese_block: SESE block to outline from the code. Parent attribute must be valid.
             :return: Replaces sese_block with a function call to the newly created function. 
        """
        target = AstToIR.deepcopySubtree(sese_block,)
        self.main_ast = getFileNode(sese_block)
        # AstToIR.link_subtree_parents(target)
        # Get the variables declared outside this scope 
        outside_scope_vars = self._getUndeclVars(target)
        # Get the declarations of those variables
        decls = []
        if (len(outside_scope_vars) > 0):
            decls = self._getDeclList(outside_scope_vars)
        if len(self._extra_decls):
            decls = list(set(decls).union(set(self._extra_decls)))
        # Build the function
        [self.outlined_subtree, self.prototype_subtree] = self.buildOutlineFunction(decls, target)
        # Build the function call
        self._funccall = self.buildOutlineCall(decls, sese_block)
        # Build list of required typedefs
        self._typedefList = self.buildFunctionTypedefList(decls, sese_block)
        # Build list of function calls inside Kernel
        self._funcCallList = self._getFuncCallList(sese_block)
        # Replace the old compound stmt with a function call
        if self._replace:
            ReplaceTool(new_node = self._funccall, old_node = sese_block).apply(sese_block.parent, getContainerAttribute(sese_block))
 

def inline_called_functions(non_visited_funcalls, objetive, 
                    native_list, allow_native):
        """ Inline called functions 
        """
        undeclared_list = []
        visited = []
        called_list = []
        while len(non_visited_funcalls) > 0:
            try:
                elem = non_visited_funcalls.pop()
                if not allow_native or (allow_native and not elem.name in native_list):
                    print " Extracting function named: " + str(elem.name)
                    funcDef = FuncDefOfNameFilter  (name = elem.name).apply(getFileNode(objetive))
                    called_list.append(funcDef) 
                    visited.append(elem.name)
                    # Check for nested calls
                    for elem in FuncCallFilter().iterate(funcDef.body):
                        if not elem in visited:
                            non_visited_funcalls.append(elem)
                elif allow_native and elem.name in native_list:
                    print " Native call to :" + str(elem.name)
                else:
                    print " Function " + elem.name + " is neither native or declared "
                    raise NotImplemented
            except NodeNotFound:
                # Declaration could not be found
                undeclared_list.append(elem)
                continue
        if len(undeclared_list) > 0:
            print " Following functions were called but not declared: "
            print [str(elem.name) for elem in undeclared_list]
	    print " Will assume they are native "
	# Reverse list to match declaration order
        called_list.reverse()
        return called_list
        

def c_outline_loop_stmts(objetive, om):
    """
        Outline a C statement from a For Loop
        
    """
    # Abort if not a SESE block   
    assert block_is_SESE(objetive)
    
    
    c_template = """
    %for typedef in typedeflist:
        ${typedef};
    %endfor
    
    %for func in funccalllist:
        ${func.decl};
    %endfor
    
       
    ${kernelProto};
    
    ${kernelFunction};
    
    
    """
    typedeflist = om.getTypedefList()

    
    allow_native = True
    native_list =  ['sqrt', 'min', 'max', 'log', 'fabs', 'pow',]
    non_visited_funcalls = om.funcCalls[:]

    called_list = inline_called_functions(non_visited_funcalls, objetive, 
		                               native_list, allow_native)
    outlinedProto = om.getOutlinedPrototype()    
    outlinedFunction = om.getOutlinedFunction()

    # Remove attributes
    for f in called_list:
        setattr(f.decl, "storage",  ['extern',])

    #===========================================================================
    from mako.template import Template
    kernel_source = Template(c_template).render_unicode(typedeflist = om.getTypedefList(),
                                                             funccalllist = called_list,
                                                             outlinedProto = outlinedProto,
                                                             kernelProto = outlinedProto,
                                                             kernelFunction = outlinedFunction,
                                                            )    
    # Restore attributes
    for f in called_list:
        # static cannot be applied anymore to that method, as it would be linked later
        setattr(f.decl, "storage",  [])


 
    return kernel_source




